from ultralytics import YOLO
import numpy as np
from typing import Optional
import queue

SEG_MODEL_PATH = "./merge_seg_rknn_model"
POSE_MODEL_PATH = "./merge_pose_rknn_model"
SEG_CORE = 0
POSE_CORE = 1

class ImageData:
    """图像数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    color_path: Optional[str] = None

class SegPose:
    def __init__(self):
        self.pose_model = YOLO(POSE_MODEL_PATH,task="pose")
        self.seg_model = YOLO(SEG_MODEL_PATH,task="segment")
        self.task_queue = queue.Queue(maxsize=10)
        self.result_queue = queue.Queue(maxsize=10)
    def put(self,image_data):
        self.task_queue.put(image_data)
        self.run()

    def run(self):
        task = self.task_queue.get()
        color_img = task.color_img
        # seg_result = self.model(task)
        seg_result = self.seg_model(color_img)
        pose_result = self.pose_model(color_img)
        
        seg_result = self._process_result(seg_result,task_type="segment")
        pose_result = self._process_result(pose_result,task_type="pose")

        result ={
            "seg_result":seg_result,
            "pose_result":pose_result,
            "image_data": task
        }

        self.result_queue.put(result)
        
    def _process_result(self, result,task_type):
        """快速结果处理"""
        try:
            if task_type == "segment":
                if not hasattr(result, 'masks') or result.masks is None:
                    return None
                masks_xy = result.masks.xy
                if len(masks_xy) == 0:
                    return None
                return masks_xy[0]
            else:  # pose
                if not hasattr(result, 'keypoints') or result.keypoints is None:
                    return None
                return result.keypoints
        except Exception as e:
            return None
        
    def get_result(self):
        return self.result_queue.get()